import json
import time
from src.core import Browser, BrowserPage

from src.core.components import NetworkData


class DeepSeekChat:
    def __init__(self, browser: Browser, page: BrowserPage):
        self.browser = browser
        self.page = page
        self.response_chunks = []
        self.completion_url = 'https://chat.deepseek.com/api/v0/chat/completion'
        self.last_chat_index = 0

    def create_new(self):
        """创建新的对话"""
        try:
            create_new_chat = self.page.page_context.get_by_text('开启新对话')
            create_new_chat.nth(1).click()
            time.sleep(1)
            print("成功创建新对话")
        except Exception as e:
            print(f"创建新对话失败: {e}")

    def send_message(self, message: str) -> str:
        """发送消息并等待响应"""
        self.response_chunks = []

        current_index = self.last_chat_index
        self.last_chat_index += 1

        # 1. 发送消息
        print(f"正在发送消息: {message[:30]}..." if len(message) > 30 else f"正在发送消息: {message}")
        chat_input = self.page.page_context.locator(selector='#chat-input')
        chat_input.fill(message)
        time.sleep(1)

        # 2. 实际发送
        send_button = self.page.page_context.locator(selector='[role=button]').nth(2)
        send_button.click()
        time.sleep(1)

        network_datas = self.page.watch(self.completion_url)
        network_data = network_datas[current_index]
        print('network_data', network_data)

        content = ''
        for stream in network_data.response_event_stream:
            if stream.get('data') == '[DONE]':
                break

            data = json.loads(stream.get('data'))
            if "stop" == data.get('choices')[0].get('finish_reason'):
                break

            if data.get('choices')[0].get('delta').get('role') == 'assistant':
                continue

            content += data.get('choices')[0].get('delta').get('content')
            pass

        return content


class DeepSeek:
    def __init__(self):
        self.browser = None
        self.page = None

    def test_mode(self):
        self.browser = Browser(headless=False)
        self.page = self.browser.open('https://chat.deepseek.com/a/chat/s/ef5caa89-d3cd-4b2e-bcd0-31d1ac149e2e')
        print("浏览器已启动，进入测试模式")

    def wait_for_close(self):
        try:
            self.browser.context.wait_for_event('close', timeout=0)
        except Exception as e:
            print(f"等待浏览器关闭时出错: {e}")

    def headless_mode(self):
        self.browser = Browser(headless=True)
        self.page = self.browser.open('https://chat.deepseek.com')
        print("浏览器已启动，进入无头模式")

        # 检测登录状态
        user_token = self.page.local_storage.get_json_storage('userToken')
        if user_token is None or user_token.get('value') is None:
            print("检测到未登录状态，切换到有界面模式以进行登录")
            self.browser.set_headless(False)
            self.page = self.browser.open('https://chat.deepseek.com')

            # 等待用户登录
            login_timeout = 300
            login_start_time = time.time()
            print(f"请在 {login_timeout // 60} 分钟内完成登录...")

            while time.time() - login_start_time < login_timeout:
                user_token = self.page.local_storage.get_json_storage('userToken')
                if user_token and user_token.get('value') is not None:
                    print("检测到成功登录")
                    break
                time.sleep(5)

            if user_token is None or user_token.get('value') is None:
                raise Exception("登录超时，请重试")

            # 登录完成后切回无头模式
            print("登录完成，切换回无头模式")
            self.browser.set_headless(True)
            self.page = self.browser.open('https://chat.deepseek.com')
        else:
            print("检测到已登录状态")

    def create_new_chat(self):
        chat = DeepSeekChat(
            browser=self.browser,
            page=self.page,
        )

        chat.create_new()

        return chat
